import math
import numpy as np
from copy import deepcopy

from History import *
from ValidityInterval import *

modelID = 0

def interpolateModels(m1, m2, state, nextState, dt):
    newModel = Model2()
    # 
    return newModel

class Model3:
    def __init__(self):
        global modelID
        self.k1 = 0
        self.k2 = 0
        self.validityInterval = ValidityInterval(4) # Number of dimension in a state
        self.goodStates = []
        self.c1GTc2 = False
        self.id = modelID
        modelID = modelID + 1
        
    def limitControls(self, c1, c2):
        if c1 < -1:
            c1 = -1
        elif c1 > 10:
            c1 = 10
        if c2 < -1:
            c2 = -1
        elif c2 > 10:
            c2 = 10
        return (c1, c2)
     
    def getControls(self, dt, a, aDot, aNew, aDotNew, stiffness):
        if self.c1GTc2:
            c2 = 0.5*((aDotNew - aDot)/dt - self.k2*math.cos(a)) / self.k1 + stiffness
            c1 = 3.0*c2 - 2.0*stiffness
        else:
            c1 = 0.5*(self.k2*math.cos(a) - (aDotNew - aDot)/dt) / self.k1 + stiffness
            c2 = 3.0*c1 - 2.0*stiffness
        #print c1, c2
        return self.limitControls(c1, c2)
    
    def setParameters(self, events):
        aRows = []
        RHSs = []
        ep = events[0]
        if ep.state.arm1.c1 > ep.state.arm1.c2:
            self.c1GTc2 = True
        else:
            self.c1GTc2 = False
        for e in events[1:]:
            if (ep.state.arm1.c1 > ep.state.arm1.c1) and not self.c1GT.c2:
                print "Unable to set model parameters"
                return
            if (ep.state.arm1.c1 < ep.state.arm1.c1) and self.c1GT.c2:
                print "Unable to set model parameters"
                return
            dt = e.time - ep.time
            RHSs.append(e.state.arm1.aDot - ep.state.arm1.aDot)
            a1 = (ep.state.arm1.c1 - ep.state.arm1.c2) * dt
            a2 = math.cos(ep.state.arm1.a) * dt
            aRows.append([a1, a2])
            ep = e
        #print RHSs
        #print aRows
        rhs = np.array(RHSs)
        a = np.array(aRows)
        x = np.linalg.lstsq(a, rhs)
        self.k1 = x[0][0]
        self.k2 = x[0][1]
        print "Model: Set parameters {0}".format(self)
        # Check model getControls
        self.testModel(events[0], events[1])
    
    def testModel(self, e1, e2):
        dt = e2.time - e1.time
        state1 = e1.state
        state2 = e2.state
        stiffness = state1.arm1.stiffness()
        (c1, c2) = self.getControls(dt, state1.arm1.a, state1.arm1.aDot, state2.arm1.a, state2.arm1.aDot, stiffness)
        (dc1, dc2) = (state1.arm1.c1-c1, state1.arm1.c2-c2)
        print "Model3: Test model ({0}, {1}), ({2}, {3}), ({4}, {5})".format(state1.arm1.c1, state1.arm1.c2, c1, c2, dc1, dc2)
        if dc1*dc1 + dc2*dc2 > 1:
            print "Model3: Model FAIL!!!!!!"

    def setGoodStates(self, events):
        ep = events[0]
        for e in events[1:]:
            # Do not use states where controls are out of bounds
            if ep.state.arm1.c1 < -0.99  or ep.state.arm1.c2 < -0.99:
                continue
            # Is ep already in good states?
            found = False
            for state in self.goodStates:
                if stateDistance(ep.state, state) < 0.001:
                    #print "Model {0}: Good state already present: {1}".format(self.id, ep.state)
                    found = True
                    break
            if found:
                ep = e
                continue
            a = ep.state.arm1.a
            aDot = ep.state.arm1.aDot
            dt = e.time - ep.time
            # Predict next state
            nextState = self.nextState(ep.state, dt)
            # print distance
            #print stateDistanceWithoutControls(nextState, e.state)
            if stateDistanceWithoutControls(nextState, e.state) / dt < 1.0:
                print "Model {0}: Found good state: {1}".format(self.id, ep.state)
                ep.state.modelExists = True
                self.addGoodState(ep.state)
            ep = e
    
    def nextState(self, state, dt):
        nextState = State()
        # Accleration
        p = self.k1 * (state.arm1.c1 - state.arm1.c2) + self.k2*math.cos(state.arm1.a)
        aDotNew = p*dt + state.arm1.aDot
        aNew = 0.5*p*dt*dt + state.arm1.aDot*dt + state.arm1.a
        nextState.arm1.a = aNew
        nextState.arm1.aDot = aDotNew
        return nextState
    
    def minDistanceToGoodState(self, state, nextState, dt, stiffness):
        dGood = 1000000
        s = deepcopy(state)
        (s.arm1.c1, s.arm1.c2) = self.getControls(dt, s.arm1.a, s.arm1.aDot, nextState.arm1.a, nextState.arm1.aDot, stiffness)
        for goodState in self.goodStates:
            #print "Model {0}: Good: {1}".format(self.id, goodState)
            dN = stateDistance(s, goodState)
            if dN < dGood:
                dGood = dN
        return dGood
    
    def addGoodState(self, state):
        self.goodStates.append(deepcopy(state))

    def __str__(self):
        return "k1: {0: 7.3f}, k2: {1: 7.3f}".format(self.k1, self.k2)

    def printMe(self):
        print "Model: k1 = " + str(self.k1)
        print "Model: k2 = " + str(self.k2)
        print "Model {0}: Good states:".format(self.id)
        for s in self.goodStates:
            print "Model:   {0}".format(s)
